import torch
import torch.nn.functional as F


class SingleLSTMNet(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, n_layer, drop_out):
        super(SingleLSTMNet, self).__init__()
        self.lstm = torch.nn.LSTM(input_size=input_dim, hidden_size=hidden_dim, num_layers=n_layer, batch_first=True,
                                  dropout=drop_out)
        self.dense = torch.nn.Linear(in_features=hidden_dim, out_features=1)

    def forward(self, x):
        out, _ = self.lstm(x)
        out = out[:, -1, :]
        out = F.relu(self.dense(out))
        return out


class NNet(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, n_layer, drop_out):
        super(NNet, self).__init__()
        self.dense1 = torch.nn.Linear(in_features=input_dim, out_features=hidden_dim)
        self.dense2 = torch.nn.Linear(in_features=hidden_dim, out_features=1)

    def forward(self, x):
        in_x = x[:, -1, :]
        hidden1 = self.dense1(in_x)
        hidden1 = F.relu(hidden1)
        out = self.dense2(hidden1)
        out = F.relu(out)
        return out
